#include "ntddk.h"
#include "Util.h"
#include "vm.h"
#include "Vmm.h"
#include "ept.h"
#include "AsmCall.h"
#include "common.h"
#include "SuperDeleteFile.h"

#define INPUT_HOOK_COMMAND CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define INPUT_UNHOOK_COMMAND CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define INPUT_Add_PID CTL_CODE(FILE_DEVICE_UNKNOWN, 0x802, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define INPUT_delete_PID CTL_CODE(FILE_DEVICE_UNKNOWN, 0x803, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define INPUT_SUPER_KILL_FILE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x804, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)

#define INPUT_HOOK_USER_PROC_ADDRESS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x805, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define INPUT_UNHOOK_USER_PROC_ADDRESS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x806, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)

PVOID pFunctionOpenProcess;
PVOID pFunctionIofCallDriver;

extern PDEVICE_OBJECT pRealDeviceObj[26];
extern PIRP pHookIrp;
extern ULONG nCurrentHookDeviceObj;

NTSTATUS VmpAsmUnHookUserProcAddress(void* context, ULONG_PTR* pParam);

void CreateProcessNotifyRoutine(_In_ HANDLE ParentId, HANDLE ProcessId, BOOLEAN Create) {
	if (!Create)
	{
		ULONG nNumberOfTotalProcessor = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
		ULONG_PTR nProcessId = (ULONG_PTR)ProcessId;
		for (ULONG processIndex = 0; processIndex < nNumberOfTotalProcessor; processIndex++)
		{
			PROCESSOR_NUMBER dtProcessorNumber = {};
			NTSTATUS nStatus = KeGetProcessorNumberFromIndex(processIndex, &dtProcessorNumber);
			if (!NT_SUCCESS(nStatus))
			{
				return;
			}

			GROUP_AFFINITY dtAffinity = {};
			dtAffinity.Group = dtProcessorNumber.Group;
			dtAffinity.Mask = 1ull << dtProcessorNumber.Number;

			GROUP_AFFINITY dtOldAffinity = {};
			KeSetSystemGroupAffinityThread(&dtAffinity, &dtOldAffinity);

			nStatus = VmpAsmUnHookUserProcAddress(0, &nProcessId);

			KeRevertToUserGroupAffinityThread(&dtOldAffinity);

		}
	}
}

static void DriverpDriverUnload(
	PDRIVER_OBJECT driver_object) {
	//UNREFERENCED_PARAMETER(driver_object);


	//HYPERPLATFORM_COMMON_DBG_BREAK();
	PsSetCreateProcessNotifyRoutine(CreateProcessNotifyRoutine, TRUE);
	VmTermination();
	//HotplugCallbackTermination();
	//PowerCallbackTermination();
	//UtilTermination();
	//PerfTermination();
	//GlobalObjectTermination();
	//LogTermination();
	UNICODE_STRING pSymbolicName;
	RtlInitUnicodeString(&pSymbolicName, L"\\??\\Whisper_Waiter");
	IoDeleteSymbolicLink(&pSymbolicName);

	for (PDEVICE_OBJECT pCurrentDeviceObj = driver_object->DeviceObject; pCurrentDeviceObj != 0;)
	{
		PDEVICE_OBJECT pTempNextDeviceObj = pCurrentDeviceObj->NextDevice;
		IoDeleteDevice(pCurrentDeviceObj);
		pCurrentDeviceObj = pTempNextDeviceObj;
	}
}

NTSTATUS VmpAsmHookUserProcAddress(void* context, ULONG_PTR *pParam) {
	KdPrint(("[Whisper]Input Commond: HUPA %p %p\n", pParam[0], pParam[1]));
	//Asm_cpuidByParam((ULONG_PTR)'HUPA', pParam);
	Asm_VmCall(HypercallNumber::kAddHookUserProcID, pParam);
	return STATUS_SUCCESS;
}

NTSTATUS VmpAsmUnHookUserProcAddress(void* context, ULONG_PTR* pParam) {
	KdPrint(("[Whisper]Input Commond: UHUP %p %p\n", pParam[0], pParam[1]));
	//Asm_cpuidByParam((ULONG_PTR)'UHUP', pParam);
	Asm_VmCall(HypercallNumber::kDeleteHookUserProcID, pParam);
	return STATUS_SUCCESS;
}

NTSTATUS VmpAsmAddPidCpuid(void* context, ULONG_PTR nPid) {
	ULONG_PTR szArrary[5];
	szArrary[0] = nPid;
	KdPrint(("[Whisper]Asm Pid!\n"));
	KdPrint(("[Whisper]Input Commond: ADPD %x\n", nPid));
	//Asm_cpuidByParam((ULONG_PTR)'ADPD', szArrary);
	Asm_VmCall(HypercallNumber::kAddAntiOpenProcID, &szArrary);
	return STATUS_SUCCESS;
}

NTSTATUS VmpAsmDePidCpuid(void* context, ULONG_PTR nPid) {
	ULONG_PTR szArrary[5];
	szArrary[0] = nPid;
	KdPrint(("[Whisper]Input Commond: DEPD %x\n", nPid));
	//Asm_cpuidByParam((ULONG_PTR)'DEPD', szArrary);
	Asm_VmCall(HypercallNumber::kDeleteAntiOpenProcID, &szArrary);
	return STATUS_SUCCESS;
}

NTSTATUS VmpAsmHookNtOpenProcess(void *context) {
	ULONG32 szArrary[5];
	//Asm_cpuid('HOOK', szArrary);
	ProcessorData* processor_data = nullptr;
	Asm_VmCall(HypercallNumber::kHookNtOpenProcess, &processor_data);
	return STATUS_SUCCESS;
}

NTSTATUS VmpAsmUnHookNtOpenProcess(void* context) {
	ULONG32 szArrary[5];
	//Asm_cpuid('UHOK', szArrary);
	ProcessorData* processor_data = nullptr;
	Asm_VmCall(HypercallNumber::kUnHookNtOpenProcess, &processor_data);
	return STATUS_SUCCESS;
}

NTSTATUS DriverNormalDeallDispatch(struct _DEVICE_OBJECT* DeviceObject, struct _IRP* Irp) {
	PIO_STACK_LOCATION pIoStackLocation = IoGetCurrentIrpStackLocation(Irp);
	Irp->IoStatus.Information = STATUS_SUCCESS;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS DriverIoControlDeallDispatch(struct _DEVICE_OBJECT* DeviceObject, struct _IRP* Irp) {
	PIO_STACK_LOCATION pIoStackLocation = IoGetCurrentIrpStackLocation(Irp);
	
	if (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode == INPUT_HOOK_COMMAND)
	{
		UtilForEachProcessor(VmpAsmHookNtOpenProcess, 0);
	} else if (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode == INPUT_UNHOOK_COMMAND)
	{
		UtilForEachProcessor(VmpAsmUnHookNtOpenProcess, 0);
	}
	else if (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode == INPUT_Add_PID)
	{
		KdPrint(("[Whisper]Add Pid!\n"));
		ULONG nNumberOfTotalProcessor = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
		for (ULONG processIndex = 0; processIndex < nNumberOfTotalProcessor; processIndex++)
		{
			PROCESSOR_NUMBER dtProcessorNumber = {};
			NTSTATUS nStatus = KeGetProcessorNumberFromIndex(processIndex, &dtProcessorNumber);
			if (!NT_SUCCESS(nStatus))
			{
				return nStatus;
			}

			GROUP_AFFINITY dtAffinity = {};
			dtAffinity.Group = dtProcessorNumber.Group;
			dtAffinity.Mask = 1ull << dtProcessorNumber.Number;

			GROUP_AFFINITY dtOldAffinity = {};
			KeSetSystemGroupAffinityThread(&dtAffinity, &dtOldAffinity);

			nStatus = VmpAsmAddPidCpuid(0, *(ULONG_PTR *)(Irp->AssociatedIrp.SystemBuffer));

			KeRevertToUserGroupAffinityThread(&dtOldAffinity);

		}
	}
	else if (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode == INPUT_delete_PID)
	{
		KdPrint(("[Whisper]Delete Pid!\n"));
		ULONG nNumberOfTotalProcessor = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
		for (ULONG processIndex = 0; processIndex < nNumberOfTotalProcessor; processIndex++)
		{
			PROCESSOR_NUMBER dtProcessorNumber = {};
			NTSTATUS nStatus = KeGetProcessorNumberFromIndex(processIndex, &dtProcessorNumber);
			if (!NT_SUCCESS(nStatus))
			{
				return nStatus;
			}

			GROUP_AFFINITY dtAffinity = {};
			dtAffinity.Group = dtProcessorNumber.Group;
			dtAffinity.Mask = 1ull << dtProcessorNumber.Number;

			GROUP_AFFINITY dtOldAffinity = {};
			KeSetSystemGroupAffinityThread(&dtAffinity, &dtOldAffinity);

			nStatus = VmpAsmDePidCpuid(0, *(ULONG_PTR*)(Irp->AssociatedIrp.SystemBuffer));

			KeRevertToUserGroupAffinityThread(&dtOldAffinity);

		}
	}
	else if (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode == INPUT_SUPER_KILL_FILE)
	{
		KdPrint(("[Whisper]Try to copy file path!\n"));
		ULONG cbin = pIoStackLocation->Parameters.DeviceIoControl.InputBufferLength;
		PCWSTR pFilePath = (PCWSTR)ExAllocatePoolWithTag(NonPagedPool, cbin + 8, kHyperPlatformCommonPoolTag);
		if (pFilePath != 0)
		{
			RtlZeroMemory((void *)pFilePath, cbin + 8);
			RtlCopyMemory((void *)pFilePath, Irp->AssociatedIrp.SystemBuffer, cbin);
			if (SuperKillFile(pFilePath))
			{
				KdPrint(("[Whisper]Kill file \'%ws\' MYABE successfully!\n", pFilePath));
			}
			else {
				KdPrint(("[Whisper]Kill file \'%ws\' MYABE failed!\n", pFilePath));
			}

			ExFreePoolWithTag((PVOID)pFilePath, kHyperPlatformCommonPoolTag);
		}
		else {
			KdPrint(("[Whisper]Allocate memory error!\n"));
		}
	}
	else if (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode == INPUT_HOOK_USER_PROC_ADDRESS)
	{
		KdPrint(("[Whisper]Try to Hook User proc address!\n"));
		ULONG cbin = pIoStackLocation->Parameters.DeviceIoControl.InputBufferLength;
		if (cbin == 0x18)
		{
			KdPrint(("[Whisper]Param %I64x %I64x!\n", (ULONG_PTR)((ULONG_PTR*)(Irp->AssociatedIrp.SystemBuffer))[0], (ULONG_PTR)((ULONG_PTR*)(Irp->AssociatedIrp.SystemBuffer))[1]));
			OBJECT_ATTRIBUTES ObjectAttributes;
			InitializeObjectAttributes(&ObjectAttributes, 0, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, 0, 0);
			CLIENT_ID clientid;
			clientid.UniqueProcess = (HANDLE)*(ULONG_PTR *)(Irp->AssociatedIrp.SystemBuffer);
			clientid.UniqueThread = 0;
			HANDLE hProcess = 0;
			NTSTATUS Status = ZwOpenProcess(&hProcess, PROCESS_ALL_ACCESS, &ObjectAttributes, &clientid);
			if (NT_SUCCESS(Status))
			{
				ULONG_PTR nParamArrary[4];
				nParamArrary[0] = (ULONG_PTR)hProcess;
				nParamArrary[1] = (ULONG_PTR)((ULONG_PTR*)(Irp->AssociatedIrp.SystemBuffer))[1];
				nParamArrary[2] = (ULONG_PTR)((ULONG_PTR*)(Irp->AssociatedIrp.SystemBuffer))[2];
				nParamArrary[3] = (ULONG_PTR)((ULONG_PTR*)(Irp->AssociatedIrp.SystemBuffer))[0];
				ULONG nNumberOfTotalProcessor = KeQueryActiveProcessorCountEx(ALL_PROCESSOR_GROUPS);
				for (ULONG processIndex = 0; processIndex < nNumberOfTotalProcessor; processIndex++)
				{
					PROCESSOR_NUMBER dtProcessorNumber = {};
					NTSTATUS nStatus = KeGetProcessorNumberFromIndex(processIndex, &dtProcessorNumber);
					if (!NT_SUCCESS(nStatus))
					{
						return nStatus;
					}

					GROUP_AFFINITY dtAffinity = {};
					dtAffinity.Group = dtProcessorNumber.Group;
					dtAffinity.Mask = 1ull << dtProcessorNumber.Number;

					GROUP_AFFINITY dtOldAffinity = {};
					KeSetSystemGroupAffinityThread(&dtAffinity, &dtOldAffinity);

					nStatus = VmpAsmHookUserProcAddress(0, (ULONG_PTR*)(nParamArrary));

					KeRevertToUserGroupAffinityThread(&dtOldAffinity);

				}

				ZwClose(hProcess);
			}
			else {
				KdPrint(("[Whisper]Open Process Error!\n"));
			}
		}
		else {
			KdPrint(("[Whisper]Buffer Length error!\n"));
		}
	}
	else if (pIoStackLocation->Parameters.DeviceIoControl.IoControlCode == INPUT_UNHOOK_USER_PROC_ADDRESS)
	{
		/*KdPrint(("[Whisper]Try to copy file path!\n"));
		ULONG cbin = pIoStackLocation->Parameters.DeviceIoControl.InputBufferLength;
		PCWSTR pFilePath = (PCWSTR)ExAllocatePoolWithTag(NonPagedPool, cbin + 8, kHyperPlatformCommonPoolTag);
		if (pFilePath != 0)
		{
			RtlZeroMemory((void*)pFilePath, cbin + 8);
			RtlCopyMemory((void*)pFilePath, Irp->AssociatedIrp.SystemBuffer, cbin);
			if (SuperKillFile(pFilePath))
			{
				KdPrint(("[Whisper]Kill file \'%ws\' MYABE successfully!\n", pFilePath));
			}
			else {
				KdPrint(("[Whisper]Kill file \'%ws\' MYABE failed!\n", pFilePath));
			}

			ExFreePoolWithTag((PVOID)pFilePath, kHyperPlatformCommonPoolTag);
		}
		else {
			KdPrint(("[Whisper]Allocate memory error!\n"));
		}*/
	}
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS DriverEntry(PDRIVER_OBJECT driver_object,
	PUNICODE_STRING registry_path) {
	PHYSICAL_ADDRESS dtPhysicalAddress = {-1};

	UNICODE_STRING dtUnicodeString;
	RtlInitUnicodeString(&dtUnicodeString, L"NtOpenProcess");
	pFunctionOpenProcess = MmGetSystemRoutineAddress(&dtUnicodeString);

	RtlInitUnicodeString(&dtUnicodeString, L"IofCallDriver");
	pFunctionIofCallDriver = MmGetSystemRoutineAddress(&dtUnicodeString);

	NTSTATUS status = STATUS_UNSUCCESSFUL;
	driver_object->DriverUnload = DriverpDriverUnload;

	KdPrint(("[Whisper]Get Pte address: %p\n", UtilGetGlobalPteAddress()));

	status = UtilInitialization(driver_object);
	if (!NT_SUCCESS(status)) {
		return status;
	}

	// Virtualize all processors
	if (VmIsWhisperInstalled())
	{
		KdPrint(("[Whisper]Whisper has been installed!\n"));
		return STATUS_SUCCESS;
	}

	if (!VmIsSupportVmx())
	{
		KdPrint(("[Whisper]System not supports Whisper!\n"));
		return STATUS_SUCCESS;
	}

	SharedProcessorData *shared_data = VmpInitializeSharedData();
	if (!shared_data)
	{
		KdPrint(("[Whisper]Initialize Shared Data Error!\n"));
		return STATUS_SUCCESS;
	}

	EptInitializeMtrrEntries();

	status = UtilForEachProcessor(VmpStartVm, shared_data);
	if (!NT_SUCCESS(status)) {
		//UtilForEachProcessor(VmpStopVm, nullptr);
		KdPrint(("[Whisper]Installed error!\n"));
		return STATUS_SUCCESS;
	}

	KdPrint(("[Whisper]Installed succeed!\n"));

	if (pFunctionIofCallDriver)
	{
		KdPrint(("[WHisper]IofCallDriver Address: 0x%p\n", pFunctionIofCallDriver));
		//FindAllDriverRealDevice();
	}

	for (size_t i = 0; i < IRP_MJ_MAXIMUM_FUNCTION; i++)
	{
		driver_object->MajorFunction[i] = DriverNormalDeallDispatch;
	}

	driver_object->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DriverIoControlDeallDispatch;

	UNICODE_STRING dtDeviceName;
	UNICODE_STRING dtSymbolicName;
	RtlInitUnicodeString(&dtDeviceName, L"\\Device\\devWhisper_Waiter");
	PDEVICE_OBJECT hDeviceObject;
	NTSTATUS ntStatus = IoCreateDevice(driver_object, (ULONG)0, &dtDeviceName, FILE_DEVICE_UNKNOWN, (ULONG)0, FALSE, &hDeviceObject);

	if (!NT_SUCCESS(ntStatus))
	{
		KdPrint(("[Whisper]Create Device failed. %X \n", ntStatus));
		return STATUS_UNSUCCESSFUL;
	}

	RtlInitUnicodeString(&dtSymbolicName, L"\\??\\Whisper_Waiter");
	ntStatus = IoCreateSymbolicLink(&dtSymbolicName, &dtDeviceName);

	if (!NT_SUCCESS(ntStatus))
	{
		KdPrint(("[Whisper]Create Symbols failed. %X \n", ntStatus));
		IoDeleteDevice(hDeviceObject);
		return STATUS_UNSUCCESSFUL;
	}

	PsSetCreateProcessNotifyRoutine(CreateProcessNotifyRoutine, FALSE);
	KdPrint(("[Whisper]start \n"));

	pHookIrp = 0;

	return STATUS_SUCCESS;
}
